/******************************************************************/
/*                                                                */
/*  Winpooch : Windows Watchdog                                   */
/*  Copyright (C) 2004-2005  Benoit Blanchon                      */
/*                                                                */
/*  This program is free software; you can redistribute it        */
/*  and/or modify it under the terms of the GNU General Public    */
/*  License as published by the Free Software Foundation; either  */
/*  version 2 of the License, or (at your option) any later       */
/*  version.                                                      */
/*                                                                */
/*  This program is distributed in the hope that it will be       */
/*  useful, but WITHOUT ANY WARRANTY; without even the implied    */
/*  warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR       */
/*  PURPOSE.  See the GNU General Public License for more         */
/*  details.                                                      */
/*                                                                */
/*  You should have received a copy of the GNU General Public     */
/*  License along with this program; if not, write to the Free    */
/*  Software Foundation, Inc.,                                    */
/*  675 Mass Ave, Cambridge, MA 02139, USA.                       */
/*                                                                */
/******************************************************************/


/******************************************************************/
/* Build configuration                                            */
/******************************************************************/

#define TRACE_LEVEL	2 /* warning */


/******************************************************************/
/* Includes                                                       */
/******************************************************************/

// module's interface
#include "ProcList.h"

// standard headers
#include <windows.h>
#include <tchar.h>
#include <tlhelp32.h>
#include <psapi.h>

// project's headers
#include "Assert.h"
#include "Strlcpy.h"
#include "Trace.h"


/******************************************************************/
/* Internal data types                                            */
/******************************************************************/

typedef struct PROCNODE {
  struct PROCNODE	*pNext ;
  PROCSTRUCT		proc ;
} PROCNODE ;

typedef struct {
  PROCNODE	*pFirst ;
  PROCNODE	*pLast ;
  HANDLE	hMutex ;
} PROCFIFO ;


/******************************************************************/
/* Internal data                                                  */
/******************************************************************/

static PROCFIFO g_fifo ;


/******************************************************************/
/* Internal functions                                             */
/******************************************************************/
	
VOID	_ProcList_Lock () ;

VOID	_ProcList_Unlock () ; 

DWORD	_ProcList_Push (PROCNODE * pNewNode) ;

BOOL	_ProcList_Pop () ;

BOOL	_ProcList_IsProcessId (DWORD dwProcessId) ;

// get process path from PID
BOOL	_ProcList_GetProcessFileName (DWORD dwProcessId, LPTSTR szFilename, DWORD nSize) ;


//PROCNODE* _ProcList_GetNode (DWORD nId)  ;


/******************************************************************/
/* Exported function : Init                                       */
/******************************************************************/

BOOL	ProcList_Init () 
{
  TRACE ;

  g_fifo.hMutex = CreateMutex (NULL, FALSE, NULL) ;
  if( ! g_fifo.hMutex )
    TRACE_ERROR (TEXT("CreateMutex failed (error=%d)\n"), GetLastError()) ;

  g_fifo.pFirst = NULL ;
  g_fifo.pLast = NULL ;

  return TRUE ;
}


/******************************************************************/
/* Exported function : Uninit                                     */
/******************************************************************/

VOID ProcList_Uninit () 
{
  TRACE ;

  while( _ProcList_Pop () ) ;

  CloseHandle (g_fifo.hMutex) ;
  g_fifo.hMutex = NULL ;
}


/******************************************************************/
/* Exported function : Add                                        */
/******************************************************************/

BOOL	ProcList_Add (const PROCSTRUCT * pProc) 
{
  PROCNODE * pNewNode ;

  TRACE_INFO ("%s (%d)\n", pProc->szName, pProc->nProcessId) ;

  pNewNode = malloc (sizeof(PROCNODE)) ;
  if( ! pNewNode ) {
    TRACE_ERROR (TEXT("malloc failed\n")) ;
    return FALSE ;
  }

  memcpy (&pNewNode->proc, pProc, sizeof(PROCSTRUCT)) ;
  
  return _ProcList_Push (pNewNode) ;
}


/******************************************************************/
/* Exported function : Enum                                       */
/******************************************************************/

VOID	ProcList_Enum (ENUMPROCCALLBACK pfnCallback, void * pContext) 
{
  PROCNODE * pNode ;

  TRACE ;

  _ProcList_Lock () ;
  
  for( pNode=g_fifo.pFirst ; pNode ; pNode=pNode->pNext )
    if( ! pfnCallback(pContext,&pNode->proc) )
      break ;

  _ProcList_Unlock () ;
}


/******************************************************************/
/* Internal function : Lock                                       */
/******************************************************************/

VOID _ProcList_Lock ()
{
  DWORD dwResult ;

  TRACE ;

  dwResult = WaitForSingleObject (g_fifo.hMutex, 2000) ;

  if( dwResult!=WAIT_OBJECT_0 )
    TRACE_ERROR (TEXT("WaitForSingleObject failed (res=%u, error=%u)\n"),
		 dwResult, GetLastError()) ;		       
}


/******************************************************************/
/* Internal function : Unlock                                     */
/******************************************************************/

VOID _ProcList_Unlock ()
{
  TRACE ;

  if( ! ReleaseMutex (g_fifo.hMutex) )
    TRACE_ERROR (TEXT("ReleaseMutex failed (error=%u)\n"), GetLastError()) ;  
}


/******************************************************************/
/* Internal function : Push                                       */
/******************************************************************/

DWORD _ProcList_Push (PROCNODE * pNewNode)
{
  TRACE ;
 
  _ProcList_Lock () ;

  // verify linked list 
  ASSERT (!g_fifo.pFirst==!g_fifo.pLast) ;

  pNewNode->pNext = NULL ;

  if( ! g_fifo.pFirst )
    g_fifo.pFirst = pNewNode ;

  if( g_fifo.pLast )
    g_fifo.pLast->pNext = pNewNode ;
  
  g_fifo.pLast = pNewNode ;

  _ProcList_Unlock () ;

  return TRUE ;
}


/******************************************************************/
/* Internal function : Pop                                        */
/******************************************************************/

BOOL _ProcList_Pop ()
{
  PROCNODE * pNode ;

  TRACE ;

  _ProcList_Lock () ;

  pNode = g_fifo.pFirst ;

  if( ! pNode ) {
    TRACE_INFO (TEXT("List is empty\n")) ;
    _ProcList_Unlock () ;
    return FALSE ;
  }
    
  g_fifo.pFirst = pNode->pNext ;

  if( !pNode->pNext ) {
    TRACE_INFO (TEXT("Remove the last node\n")) ;
    g_fifo.pLast = NULL ;    
  }

  _ProcList_Unlock () ;

  free (pNode) ;

  return TRUE ;
}


/******************************************************************/
/* Internal function : IsProcessId                                */
/******************************************************************/

BOOL	_ProcList_IsProcessId (DWORD dwProcessId)
{
  HANDLE	hProcess ;

  hProcess = OpenProcess (0, FALSE, dwProcessId) ;

  if( ! hProcess )
    return ERROR_ACCESS_DENIED==GetLastError() ;
  
  CloseHandle (hProcess) ;
  
  return TRUE ;
}


/******************************************************************/
/* Internal function : GetProcessFileName                         */
/******************************************************************/

BOOL _ProcList_GetProcessFileName (DWORD dwProcessId, LPTSTR szFilename, DWORD nSize) 
{
  HANDLE	hProcess ;
  BOOL		bSuccess ;
  DWORD		dwAccess = PROCESS_QUERY_INFORMATION|PROCESS_VM_READ ;

  ASSERT (szFilename!=NULL) ;
  ASSERT (nSize>0) ;

  if( dwProcessId==0 )
    return _tcslcpy (szFilename, TEXT("System"), nSize) ;
    
  hProcess = OpenProcess (dwAccess,FALSE,dwProcessId) ;

  if( ! hProcess ) {
    TRACE_ERROR (TEXT("OpenProcess failed (error=%d)\n"), GetLastError()) ;
    return FALSE ;
  }

  szFilename[0] = 0 ;
  bSuccess = GetModuleFileNameEx (hProcess, NULL, szFilename, nSize) ;

  if( ! bSuccess )
    TRACE_ERROR (TEXT("GetModuleFileNameEx failed (pid=%u, error=%d)\n"), 
		 dwProcessId, GetLastError()) ;

  CloseHandle (hProcess) ;

  return bSuccess ;
}


/******************************************************************/
/* Exported function : Remove                                     */
/******************************************************************/

VOID ProcList_Remove (DWORD dwProcessId)
{
  PROCNODE	*pCur, *pPrev ;

  TRACE ;
  
  _ProcList_Lock () ;

  pPrev = NULL ;
  for( pCur=g_fifo.pFirst ; pCur ; pCur=pCur->pNext )
    {
      if( pCur->proc.nProcessId==dwProcessId )
	{
	  if( pPrev ) pPrev->pNext = pCur->pNext ;
	  else g_fifo.pFirst = pCur->pNext ;
	  if( ! pCur->pNext ) g_fifo.pLast = pPrev ;
	    
	  free (pCur) ;
	  pCur = pPrev ;
	}

      pPrev = pCur ;
    }

  _ProcList_Unlock () ;
}


/******************************************************************/
/* Exported function : Discover                                   */
/******************************************************************/

INT	ProcList_Discover (ENUMPROCCALLBACK pfnAddedProcesses, void * pAddedContext, 
			   ENUMPROCCALLBACK pfnRemovedProcesses, void * pRemovedContext) 
{
  HANDLE		hSnapShot ;
  PROCNODE		*pCur, *pPrev ;
  PROCESSENTRY32	pe32 ;
  BOOL			bNewProcess ;
  UINT			nChanges = 0 ;

  TRACE ;
  
  _ProcList_Lock () ;

  //
  // Look for removed processes
  //

  pPrev = NULL ;
  for( pCur=g_fifo.pFirst ; pCur ; pCur=pCur->pNext )
    {
      if( ! _ProcList_IsProcessId (pCur->proc.nProcessId) )
	{
	  TRACE_INFO (TEXT("Process %d removed\n"), pCur->proc.nProcessId) ;
	  nChanges++ ;

	  if( pPrev ) pPrev->pNext = pCur->pNext ;
	  else g_fifo.pFirst = pCur->pNext ;
	  if( ! pCur->pNext ) g_fifo.pLast = pPrev ;
	    
	  if( pfnRemovedProcesses )
	    pfnRemovedProcesses (pRemovedContext, &pCur->proc) ;

	  free (pCur) ;
	  pCur = pPrev ;
	}

      pPrev = pCur ;
    }


  //
  // Look for new processes 
  //

  hSnapShot = CreateToolhelp32Snapshot (TH32CS_SNAPPROCESS, 0) ;
  if( hSnapShot==INVALID_HANDLE_VALUE ) {
    TRACE_ERROR (TEXT("CreateToolhelp32Snapshot failed (error=%d)\n"), GetLastError()) ;
    _ProcList_Unlock () ;
    return -1 ;
  }
  
  pe32.dwSize = sizeof(pe32) ;

  if( ! Process32First (hSnapShot, &pe32 ) )
    {
      TRACE_ERROR ("Process32First failed (error=%d)\n", GetLastError()) ;
      CloseHandle (hSnapShot) ;     
      _ProcList_Unlock () ; 
      return -1 ;
    } 

  do {

    // skip process 0
    if( pe32.th32ProcessID==0 ) continue ;

    // is process in list ?
    bNewProcess = TRUE ;
    for( pCur=g_fifo.pFirst ; pCur&&bNewProcess ; pCur=pCur->pNext )
      if( pCur->proc.nProcessId ==  pe32.th32ProcessID )
	bNewProcess = FALSE ;

    if( bNewProcess )
      {
	TRACE_INFO (TEXT("Process %d added\n"), pe32.th32ProcessID) ;
	nChanges++ ;

	pCur = malloc (sizeof(PROCNODE)) ;
	
	pCur->proc.nProcessId = pe32.th32ProcessID ;
	_tcslcpy (pCur->proc.szName, pe32.szExeFile, 32) ;
	pCur->proc.nState = PS_UNKNOWN_STATE ;
	pCur->proc.szPath[0] = 0 ;
	_ProcList_GetProcessFileName (pe32.th32ProcessID, pCur->proc.szPath, MAX_PATH) ;

	pCur->pNext = NULL ;

	if( ! g_fifo.pFirst ) g_fifo.pFirst = pCur ;
	if( g_fifo.pLast ) g_fifo.pLast->pNext = pCur ;
	g_fifo.pLast = pCur ;

	if( pfnAddedProcesses )
	  pfnAddedProcesses (pAddedContext, &pCur->proc) ;
      } 
  }
  while( Process32Next(hSnapShot,&pe32) ) ;
  
  CloseHandle (hSnapShot) ;

  _ProcList_Unlock () ;

  TRACE ;
  
  return nChanges ;
  
}


/******************************************************************/
/* Exported function : Get                                        */
/******************************************************************/

PROCSTRUCT * ProcList_Get (DWORD dwProcessId)
{
  PROCNODE		*pCur ;

  TRACE ;
  
  _ProcList_Lock () ;

  for( pCur=g_fifo.pFirst ; pCur ; pCur=pCur->pNext )
    if( pCur->proc.nProcessId == dwProcessId )
      break ;
    
  if( ! pCur ) { 
    _ProcList_Unlock () ;
    return NULL ;
  }

  // the unlock is in ProcList_Release

  return &pCur->proc ;
}


/******************************************************************/
/* Exported function : Release                                    */
/******************************************************************/

VOID	ProcList_Release (PROCSTRUCT * pProc) 
{
  TRACE ;

  if( ! pProc ) return ;

  _ProcList_Unlock () ;
}


/******************************************************************/
/* Exported function : IsHooked                                   */
/******************************************************************/

BOOL	ProcList_IsHooked (PROCSTRUCT * pProc) 
{
  // verify params
  ASSERT (pProc!=NULL) ;

  return pProc->nState==PS_HOOKED_SINCE_BIRTH
    ||   pProc->nState==PS_HOOKED_WHILE_RUNNING
    ||   pProc->nState==PS_HOOKED_BY_PREV_INSTANCE ;
}


/******************************************************************/
/* Exported function : GetState                                   */
/******************************************************************/

DWORD ProcList_GetState (DWORD dwProcessId)
{
  PROCSTRUCT	*pProc ;
  DWORD		nState ;

  TRACE ;

  pProc = ProcList_Get (dwProcessId) ;  

  if( ! pProc ) nState = PS_UNKNOWN_STATE ; 
  else nState = pProc->nState ;

  ProcList_Release (pProc) ;

  return nState ;
}


/******************************************************************/
/* Exported function : SetState                                   */
/******************************************************************/

VOID ProcList_SetState (DWORD dwProcessId, DWORD nState)
{
  PROCSTRUCT	*pProc ;

  TRACE ;

  pProc = ProcList_Get (dwProcessId) ;  

  if( pProc ) pProc->nState = nState ;

  ProcList_Release (pProc) ;
}

